#include <iostream>
#include <immintrin.h>
#include <cstring> // 用于memcpy

// 通用C++转置实现（原地转置）
void transpose_cpp(int mat[4][4]) {
    for(int i=0; i<4; ++i)
        for(int j=i+1; j<4; ++j)
            std::swap(mat[i][j], mat[j][i]);
}

// AVX2优化转置（非破坏性转置）
void transpose_avx2(const int src[4][4], int dst[4][4]) {
    // 加载4行数据到寄存器
    __m128i row0 = _mm_loadu_si128((__m128i*)&src[0]);
    __m128i row1 = _mm_loadu_si128((__m128i*)&src[1]);
    __m128i row2 = _mm_loadu_si128((__m128i*)&src[2]);
    __m128i row3 = _mm_loadu_si128((__m128i*)&src[3]);

    // 第一阶段：32位整数的解包操作
    __m128i tmp0 = _mm_unpacklo_epi32(row0, row1); // [a0 b0 a1 b1]
    __m128i tmp1 = _mm_unpacklo_epi32(row2, row3); // [c0 d0 c1 d1]
    __m128i tmp2 = _mm_unpackhi_epi32(row0, row1); // [a2 b2 a3 b3]
    __m128i tmp3 = _mm_unpackhi_epi32(row2, row3); // [c2 d2 c3 d3]

    // 第二阶段：使用浮点移动指令重组数据
    row0 = _mm_castps_si128(_mm_movelh_ps(
        _mm_castsi128_ps(tmp0), _mm_castsi128_ps(tmp1)));
    row1 = _mm_castps_si128(_mm_movehl_ps(
        _mm_castsi128_ps(tmp1), _mm_castsi128_ps(tmp0)));
    row2 = _mm_castps_si128(_mm_movelh_ps(
        _mm_castsi128_ps(tmp2), _mm_castsi128_ps(tmp3)));
    row3 = _mm_castps_si128(_mm_movehl_ps(
        _mm_castsi128_ps(tmp3), _mm_castsi128_ps(tmp2)));

    // 存储转置结果
    _mm_storeu_si128((__m128i*)&dst[0], row0);
    _mm_storeu_si128((__m128i*)&dst[1], row1);
    _mm_storeu_si128((__m128i*)&dst[2], row2);
    _mm_storeu_si128((__m128i*)&dst[3], row3);
}

// 打印矩阵
void print_matrix(const int mat[4][4]) {
    for(int i=0; i<4; ++i) {
        for(int j=0; j<4; ++j)
            std::cout << mat[i][j] << "\t";
        std::cout << "\n";
    }
}

int main() {
    // 初始化测试矩阵
    int matrix_orig[4][4] = {
        { 1,  2,  3,  4},
        { 5,  6,  7,  8},
        { 9, 10, 11, 12},
        {13, 14, 15, 16}
    };
    
    // 测试C++版本
    int matrix_cpp[4][4];
    memcpy(matrix_cpp, matrix_orig, sizeof(matrix_orig));
    
    std::cout << "Original Matrix:\n";
    print_matrix(matrix_cpp);
    
    transpose_cpp(matrix_cpp);
    std::cout << "\nC++ Transposed:\n";
    print_matrix(matrix_cpp);

    // 测试AVX2版本
    int matrix_avx[4][4];
    transpose_avx2(matrix_orig, matrix_avx);
    
    std::cout << "\nAVX2 Transposed:\n";
    print_matrix(matrix_avx);

    // 验证结果一致性
    bool match = true;
    for(int i=0; i<4; ++i)
        for(int j=0; j<4; ++j)
            if(matrix_cpp[i][j] != matrix_avx[i][j])
                match = false;
                
    std::cout << "\nResult Verification: " 
              << (match ? "PASS" : "FAIL") << std::endl;

    return 0;
}